# Copyright (c) 2014 City of Hope
# Authors: Juan Du and Dustin E Schones

"""
This module contains all the classes to deal with BED files
"""
import re, os, shutil, time, sys
from math import *
from string import *
import GenomeData;


plus = re.compile('\+');
minus = re.compile('\-');

bedError = "Error in BED class";
"""
  BED types:

  BED3: chrom, start and end
  BED6: BED3 + name + score + strand ('+' or '-')
  BEDPE: BED3 + BED3 + name + score + strand ('+' or '-') + strand ('-' or '+')
"""

#----------------------------
#----------------------------

class BPAIR:
    def __init__(self, start, end):
        self.start = start
        self.end = end

    def __set__(self, start, end):
        self.start = start
        self.end = end

    def getCoord(self):
        outstring=str(self.start) + "\t" + str(self.end)
        try:
            return outstring
        except:
            sys.stderr.write("No coord information for %s\n" % self)
            return ''
#----------------------------
#----------------------------

class BED3:
    """
    Class for bed lines with 3 values: chrom, start and end
    """
    def __init__(self, chrom, start, end):
        self.chrom = chrom;
        self.start = start;
        self.end = end;
    def __set__(self, chrom, start, end):
        self.chrom = chrom;
        self.start = start;
        self.end = end;
    def getCoord(self):
        outstring = self.chrom + "\t" + str(self.start) + "\t" + str(self.end);
        try:
            return outstring;
        except:
            sys.stderr.write("No coord information for %s\n" % self)
            return ''
#----------------------------
#----------------------------

class BED6:
    """
    Class for bed lines with 6 values:  chrom, start, end, name, score, strand
    """
    def __init__(self, chrom, start, end, name, score, strand):
        self.chrom = chrom;
        self.start = start;
        self.end = end;
        self.name = name;
        self.score = score;
        self.strand = strand;
    def __set__(self, chrom, start, end, name, score, strand):
        self.chrom = chrom;
        self.start = start;
        self.end = end;
        self.name = name;
        self.score = score;
        self.strand = strand;
    def getCoord(self):
        outstring = self.chrom + "\t" + str(self.start) + "\t" + \
                    str(self.end) + "\t" + self.name + "\t" + \
                    str(self.score) + "\t" + self.strand;
        try:
            return outstring;
        except:
            sys.stderr.write("No coord information for %s\n" % self)
            return ''



class PAIR:
    """
    Class to deal with paired-end read bed files
    """
    def __init__(self, one, two):
        self.one = one;
        self.two = two;
    def __set__(self, one, two):
        self.one = one;
        self.two = two;

#----------------------------
#----------------------------

class BED:
    """
    Class to deal with bed files and do common operations 
    """

    def __init__(self, species="hg19", file=None, bed_type="BED3", val_threshold=-1e308):

        """ Overload __init__ so that if a threshold is given, only
        grab bed vals that are above threshold -- This won't do
        anything different for bed3 entries because there is no value
        information in these.

        Reads in a bed file and builds a dictionary with chromosomes
        as keys and lists of bed elements as values
        """

        self.bed_vals = {}

        """
        initialize a dictionary with chromosomes
        """
        for c in GenomeData.species_chroms[species]:
            self.bed_vals[c] = [];

        if(file):
            if re.match(bed_type, "BED3"):	
                infile = open(file);
                for line in infile:

                    """ check to make sure not a header line """
                    if not re.match("track", line):
                        line = line.strip();
                        sline = line.split();

                        if len(sline) == 3:
                            bed = BED3(sline[0], atoi(sline[1]), atoi(sline[2]));
                            self.bed_vals[sline[0]].append(bed);

                        if len(sline) == 4:
                            if atof(sline[3]) >= val_threshold:
                                bed = BED3(sline[0], atoi(sline[1]), atoi(sline[2]));
                                self.bed_vals[sline[0]].append(bed);

                        if len(sline) == 6:
                            if atof(sline[4]) >= val_threshold:
                                bed = BED3(sline[0], atoi(sline[1]), atoi(sline[2]));
                                self.bed_vals[sline[0]].append(bed);
                        if len(sline) != 3 and len(sline) !=4 and len(sline) != 6:

                            bed = BED3(sline[0], atoi(sline[1]), atoi(sline[2]));
                            self.bed_vals[sline[0]].append(bed);


            elif re.match(bed_type, "BED6"):
                """ If want BED6 element """
                infile = open(file);
                for line in infile:
                    """ check to make sure not a header line """
                    if not re.match("track", line):
                        line = line.strip();
                        sline = line.split();
                        if len(sline) == 3:
                            sys.stderr.write("Need BED6 to make BED6")
                            raise bedError
                        if len(sline) == 4:
                            sys.stderr.write("Need BED6 to make BED6")
                            raise bedError
                        if len(sline) == 6:
                            if atof(sline[4]) >= val_threshold:
                                bed = BED6(sline[0], atoi(sline[1]), atoi(sline[2]),
                                           sline[3], atof(sline[4]), sline[5]);
                                self.bed_vals[sline[0]].append(bed);


    def keys(self):
        """
        Return a list of the keys - duplicating the function of a dictionary
        """
        return self.bed_vals.keys()


    def getStarts_consider_strands(self, chrom):
        """
        Return a list of starts on a given chromosome
        """
        starts = [];
        for t in self.bed_vals[chrom]:
            if plus.match(t.strand):
                starts.append(t.start);
            elif minus.match(t.strand):
                starts.append(t.end);
        try:
            return starts;
        except:
            sys.stderr.write("Having trouble returning starts %s\n" % self)
            return ''

    def getStarts(self, chrom):
        """
        Return a list of starts on a given chromosome
        """
        starts = [];
        for t in self.bed_vals[chrom]:
            starts.append(t.start);
        try:
            return starts;
        except:
            sys.stderr.write("Having trouble returning starts %s\n" % self)
            return ''


    def getEnds(self, chrom):
        """
        Return a list of starts on a given chromosome
        """
        ends = [];
        for t in self.bed_vals[chrom]:
            ends.append(t.end);
        try:
            return ends;
        except:
            sys.stderr.write("Having trouble returning ends %s\n" % self)
            return ''

    def getChroms(self):
        """
        Return a list of all the chromosomes in order (ordered keys)
        """
        try:
            return chroms[species];
        except:
            sys.stderr.write("Can't return chromosome list\n" % self)
            return ''


    def getNumVals(self):
        """
        Return the number of bed vals in BED instance
        """
        num = 0;
        for c in self.bed_vals.keys():
            num += len(self.bed_vals[c]);
        return num;


    def addChrom(self, chrom, bed_list):
        if self.bed_vals.has_key(chrom) and len(self.bed_vals[chrom]) > 0:
            sys.stderr.write("chromsome %s already populated\n" % chrom)
            raise bedError
        else: self.bed_vals[chrom] = bed_list;


    def __del__(self):
        """
        Delete, delete;
        """
        self.bed_vals.clear()


    def __contains__(self, item):
        """
        Returns  mapping iterator
        """
        return self.bed_vals.has_key(item)

    def __iter__(self):
        """
        Returns mapping iterator
        """
        return self.bed_vals.iterkeys()

    def __len__(self):
        """
        Returns number of bed_vals
        """
        return len(self.bed_vals)

    def __delitem__(self, name):
        """
        removes a chrom if its name exists in the dictionary
        -- I guess this could possible be useful at some point

        """
        if self.bed_vals.has_key(name):
            del self.bed_vals[name]
        else: raise bedError

    def __setitem__(self, name, bedlist):
        """
        Sets a new bed value
        """
        self.bed_vals[name] = bedlist

    def __getitem__(self, name):
        """
        Returns a bed_val indexed by its name or None if no such bed_val exists
        """
        if self.bed_vals.has_key(name):
            return self.bed_vals[name]
        else: raise bedError



class BEDPE:
    """
    Class to deal with bed files and do common operations 
    """ 
    def __init__(self, species="hg19", file=None, bed_type="BPAIR", val_threshold=-1e308):

        """ Overload __init__ so that if a threshold is given, only
        grab bed vals that are above threshold -- This won't do
        anything different for entries with no value information

        Reads in a bed file and builds a dictionary with chromosomes
        as keys and lists of bed elements as values

        Example line:
        chr7    53242933        53243032        chr7    53242939        53243038        HWI-ST667_0142:7:1101:1045:41063#GCTAAT 255     +   -



        ** RIGHT NOW ONLY DOING BED2 **
        """

        self.bed_pairs = {}
        """
        initialize a dictionary with chromosomes
        """
        for c in GenomeData.species_chroms[species]:
            self.bed_pairs[c] = [];

        if(file):


            if re.match(bed_type, "BPAIR"):
                """ store pair information - right now only storing
                BED3 information for each individual read
                IMPORTANT: This assumes that pairs have aleady been filtered
                to only keep those pairs where:
                (1) both are on the same chromosome
                (2) they are properly oriented (i.e. facing each other)
                """
                infile=open(file)
                for line in infile:
                    if not re.match("track", line):
                        line=line.strip()
                        line=line.split()
                        if plus.match(line[8]):
                            info=BPAIR(atoi(line[1]), atoi(line[5]))

                        elif minus.match(line[8]):
                            info=BPAIR(atoi(line[4]), atoi(line[2]))
                        self.bed_pairs[line[0]].append(info)

    def keys(self):
        """
        Return a list of the keys - duplicating the function of a dictionary
        """
        return self.bed_pairs.keys()

   
    def getChroms(self):
        """
        Return a list of all the chromosomes in order (ordered keys)
        """
        try:
            return chroms[species];
        except:
            sys.stderr.write("Can't return chromosome list\n" % self)
            return ''

   
    def getNumPairs(self):
        """
        Return the number of bed vals in BED instance
        """
        num = 0;
        for c in self.bed_pairs.keys():
            num += len(self.bed_pairs[c]);
        return num;


    def addChrom(self, chrom, bed_list):
        if self.bed_pairs.has_key(chrom) and len(self.bed_pairs[chrom]) > 0:
            sys.stderr.write("chromsome %s already populated\n" % chrom)
            raise bedError
        else: self.bed_pairs[chrom] = bed_list;
        

    def __del__(self):
        """
        Delete, delete;
        """
        self.bed_pairs.clear()


    def __contains__(self, item):
        """
        Returns  mapping iterator
        """
        return self.bed_pairs.has_key(item)

    def __iter__(self):
        """
        Returns mapping iterator
        """
        return self.bed_pairs.iterkeys()

    def __len__(self):
        """
        Returns number of bed_pairs
        """
        return len(self.bed_pairs)

    def __delitem__(self, name):
        """
        removes a chrom if its name exists in the dictionary -- I
        guess this could possible be useful at some point
        """
        if self.bed_pairs.has_key(name):
            del self.bed_pairs[name]
        else: raise bedError
        
    def __setitem__(self, name, bedlist):
        """
        Sets a new bed value
        """
        self.bed_pairs[name] = bedlist

    def __getitem__(self, name):
        """ Returns a bed_pair indexed by its name or None if no such
        bed_pair exists """
        if self.bed_pairs.has_key(name):
            return self.bed_pairs[name]
        else: raise bedError
